package com.flink.hbase.statemanagement.managed;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
import org.apache.flink.util.Collector;

import java.time.LocalDateTime;
import java.util.Random;

/**
 * ValueState托管状态示例
 * 
 * 展示如何使用ValueState进行状态管理：
 * 1. 用户登录状态跟踪
 * 2. 交易金额累加器
 * 3. 会话超时检测
 * 4. 状态过期清理
 * 
 * @author Flink Team
 */
public class ValueStateExample {
    
    /**
     * 用户登录事件
     */
    public static class UserLoginEvent {
        public String userId;
        public String loginTime;
        public String ipAddress;
        public String deviceType;
        
        public UserLoginEvent() {}
        
        public UserLoginEvent(String userId, String loginTime, String ipAddress, String deviceType) {
            this.userId = userId;
            this.loginTime = loginTime;
            this.ipAddress = ipAddress;
            this.deviceType = deviceType;
        }
        
        @Override
        public String toString() {
            return String.format("UserLoginEvent{userId='%s', loginTime='%s', ipAddress='%s', deviceType='%s'}", 
                userId, loginTime, ipAddress, deviceType);
        }
    }
    
    /**
     * 交易事件
     */
    public static class TransactionEvent {
        public String userId;
        public String transactionId;
        public double amount;
        public long timestamp;
        
        public TransactionEvent() {}
        
        public TransactionEvent(String userId, String transactionId, double amount, long timestamp) {
            this.userId = userId;
            this.transactionId = transactionId;
            this.amount = amount;
            this.timestamp = timestamp;
        }
        
        @Override
        public String toString() {
            return String.format("TransactionEvent{userId='%s', transactionId='%s', amount=%.2f, timestamp=%d}", 
                userId, transactionId, amount, timestamp);
        }
    }
    
    /**
     * 用户状态信息
     */
    public static class UserStateInfo {
        public String userId;
        public String lastLoginTime;
        public String lastLoginIp;
        public int loginCount;
        public double totalAmount;
        public long lastActivityTime;
        
        public UserStateInfo() {}
        
        public UserStateInfo(String userId, String lastLoginTime, String lastLoginIp, 
                           int loginCount, double totalAmount, long lastActivityTime) {
            this.userId = userId;
            this.lastLoginTime = lastLoginTime;
            this.lastLoginIp = lastLoginIp;
            this.loginCount = loginCount;
            this.totalAmount = totalAmount;
            this.lastActivityTime = lastActivityTime;
        }
        
        @Override
        public String toString() {
            return String.format("UserStateInfo{userId='%s', lastLoginTime='%s', lastLoginIp='%s', " +
                "loginCount=%d, totalAmount=%.2f, lastActivityTime=%d}", 
                userId, lastLoginTime, lastLoginIp, loginCount, totalAmount, lastActivityTime);
        }
    }
    
    /**
     * 用户登录状态跟踪函数
     * 使用ValueState跟踪用户登录信息
     */
    public static class UserLoginStateFunction extends RichFlatMapFunction<UserLoginEvent, UserStateInfo> {
        
        // 用户登录状态
        private ValueState<UserStateInfo> userStateValueState;
        
        @Override
        public void open(Configuration parameters) throws Exception {
            // 创建ValueState描述符
            ValueStateDescriptor<UserStateInfo> userStateDescriptor = new ValueStateDescriptor<>(
                "user-login-state", // 状态名称
                TypeInformation.of(UserStateInfo.class) // 状态类型
            );
            
            // 设置状态TTL（可选）
            // StateTtlConfig ttlConfig = StateTtlConfig
            //     .newBuilder(Time.hours(24))
            //     .setUpdateType(StateTtlConfig.UpdateType.OnCreateAndWrite)
            //     .setStateVisibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
            //     .build();
            // userStateDescriptor.enableTimeToLive(ttlConfig);
            
            // 获取状态
            userStateValueState = getRuntimeContext().getState(userStateDescriptor);
        }
        
        @Override
        public void flatMap(UserLoginEvent event, Collector<UserStateInfo> collector) throws Exception {
            // 获取当前用户状态
            UserStateInfo currentState = userStateValueState.value();
            
            if (currentState == null) {
                // 首次登录，创建新状态
                currentState = new UserStateInfo(
                    event.userId,
                    event.loginTime,
                    event.ipAddress,
                    1,
                    0.0,
                    System.currentTimeMillis()
                );
                System.out.println("🆕 用户首次登录: " + event.userId);
            } else {
                // 更新登录状态
                currentState.lastLoginTime = event.loginTime;
                currentState.lastLoginIp = event.ipAddress;
                currentState.loginCount += 1;
                currentState.lastActivityTime = System.currentTimeMillis();
                
                System.out.println("🔄 用户再次登录: " + event.userId + "，登录次数: " + currentState.loginCount);
            }
            
            // 更新状态
            userStateValueState.update(currentState);
            
            // 输出状态信息
            collector.collect(currentState);
        }
    }
    
    /**
     * 交易金额累加器函数
     * 使用ValueState累加用户交易金额
     */
    public static class TransactionAmountAccumulator extends KeyedProcessFunction<String, TransactionEvent, Tuple3<String, Double, Integer>> {
        
        // 累计交易金额状态
        private ValueState<Double> totalAmountState;
        // 交易次数状态
        private ValueState<Integer> transactionCountState;
        // 最后更新时间状态
        private ValueState<Long> lastUpdateTimeState;
        
        @Override
        public void open(Configuration parameters) throws Exception {
            // 创建状态描述符
            ValueStateDescriptor<Double> totalAmountDescriptor = new ValueStateDescriptor<>(
                "total-amount", Double.class, 0.0);
            
            ValueStateDescriptor<Integer> transactionCountDescriptor = new ValueStateDescriptor<>(
                "transaction-count", Integer.class, 0);
            
            ValueStateDescriptor<Long> lastUpdateTimeDescriptor = new ValueStateDescriptor<>(
                "last-update-time", Long.class, 0L);
            
            // 获取状态
            totalAmountState = getRuntimeContext().getState(totalAmountDescriptor);
            transactionCountState = getRuntimeContext().getState(transactionCountDescriptor);
            lastUpdateTimeState = getRuntimeContext().getState(lastUpdateTimeDescriptor);
        }
        
        @Override
        public void processElement(TransactionEvent event, Context context, 
                                 Collector<Tuple3<String, Double, Integer>> collector) throws Exception {
            
            // 获取当前状态值
            Double currentAmount = totalAmountState.value();
            Integer currentCount = transactionCountState.value();
            Long lastUpdateTime = lastUpdateTimeState.value();
            
            // 更新状态
            double newAmount = currentAmount + event.amount;
            int newCount = currentCount + 1;
            long currentTime = System.currentTimeMillis();
            
            // 保存更新后的状态
            totalAmountState.update(newAmount);
            transactionCountState.update(newCount);
            lastUpdateTimeState.update(currentTime);
            
            // 输出结果
            collector.collect(new Tuple3<>(event.userId, newAmount, newCount));
            
            System.out.println(String.format("💰 用户 %s 交易更新: 金额=%.2f, 次数=%d, 本次金额=%.2f", 
                event.userId, newAmount, newCount, event.amount));
            
            // 设置定时器，用于状态清理（例如：24小时后清理）
            context.timerService().registerProcessingTimeTimer(currentTime + 24 * 60 * 60 * 1000L);
        }
        
        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, 
                           Collector<Tuple3<String, Double, Integer>> out) throws Exception {
            
            // 检查状态是否过期需要清理
            Long lastUpdateTime = lastUpdateTimeState.value();
            if (lastUpdateTime != null && timestamp - lastUpdateTime > 24 * 60 * 60 * 1000L) {
                // 清理过期状态
                totalAmountState.clear();
                transactionCountState.clear();
                lastUpdateTimeState.clear();
                
                System.out.println("🧹 清理过期状态: " + ctx.getCurrentKey());
            }
        }
    }
    
    /**
     * 会话超时检测函数
     * 使用ValueState跟踪用户会话状态
     */
    public static class SessionTimeoutDetector extends KeyedProcessFunction<String, UserLoginEvent, Tuple2<String, String>> {
        
        // 会话开始时间状态
        private ValueState<Long> sessionStartTimeState;
        // 会话状态
        private ValueState<String> sessionStatusState;
        
        private static final long SESSION_TIMEOUT = 30 * 60 * 1000L; // 30分钟超时
        
        @Override
        public void open(Configuration parameters) throws Exception {
            ValueStateDescriptor<Long> sessionStartTimeDescriptor = new ValueStateDescriptor<>(
                "session-start-time", Long.class);
            
            ValueStateDescriptor<String> sessionStatusDescriptor = new ValueStateDescriptor<>(
                "session-status", String.class);
            
            sessionStartTimeState = getRuntimeContext().getState(sessionStartTimeDescriptor);
            sessionStatusState = getRuntimeContext().getState(sessionStatusDescriptor);
        }
        
        @Override
        public void processElement(UserLoginEvent event, Context context, 
                                 Collector<Tuple2<String, String>> collector) throws Exception {
            
            long currentTime = System.currentTimeMillis();
            Long sessionStartTime = sessionStartTimeState.value();
            String sessionStatus = sessionStatusState.value();
            
            if (sessionStartTime == null) {
                // 新会话开始
                sessionStartTimeState.update(currentTime);
                sessionStatusState.update("ACTIVE");
                
                collector.collect(new Tuple2<>(event.userId, "SESSION_STARTED"));
                System.out.println("🎯 会话开始: " + event.userId);
                
                // 设置会话超时定时器
                context.timerService().registerProcessingTimeTimer(currentTime + SESSION_TIMEOUT);
                
            } else if ("ACTIVE".equals(sessionStatus)) {
                // 活跃会话，延长会话时间
                sessionStartTimeState.update(currentTime);
                
                collector.collect(new Tuple2<>(event.userId, "SESSION_RENEWED"));
                System.out.println("🔄 会话续期: " + event.userId);
                
                // 重新设置超时定时器
                context.timerService().registerProcessingTimeTimer(currentTime + SESSION_TIMEOUT);
            }
        }
        
        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, 
                           Collector<Tuple2<String, String>> out) throws Exception {
            
            Long sessionStartTime = sessionStartTimeState.value();
            String sessionStatus = sessionStatusState.value();
            
            if (sessionStartTime != null && "ACTIVE".equals(sessionStatus)) {
                // 会话超时
                sessionStatusState.update("EXPIRED");
                
                out.collect(new Tuple2<>(ctx.getCurrentKey(), "SESSION_EXPIRED"));
                System.out.println("⏰ 会话超时: " + ctx.getCurrentKey());
                
                // 清理状态
                sessionStartTimeState.clear();
                sessionStatusState.clear();
            }
        }
    }
    
    /**
     * 模拟用户登录事件源
     */
    public static class MockUserLoginEventSource implements SourceFunction<UserLoginEvent> {
        private volatile boolean running = true;
        private final Random random = new Random();
        private final String[] userIds = {"user1", "user2", "user3", "user4", "user5"};
        private final String[] ipAddresses = {"192.168.1.1", "192.168.1.2", "10.0.0.1", "10.0.0.2"};
        private final String[] deviceTypes = {"mobile", "desktop", "tablet"};
        
        @Override
        public void run(SourceContext<UserLoginEvent> ctx) throws Exception {
            while (running) {
                UserLoginEvent event = new UserLoginEvent(
                    userIds[random.nextInt(userIds.length)],
                    LocalDateTime.now().toString(),
                    ipAddresses[random.nextInt(ipAddresses.length)],
                    deviceTypes[random.nextInt(deviceTypes.length)]
                );
                
                ctx.collect(event);
                Thread.sleep(2000 + random.nextInt(3000)); // 2-5秒随机间隔
            }
        }
        
        @Override
        public void cancel() {
            running = false;
        }
    }
    
    /**
     * 模拟交易事件源
     */
    public static class MockTransactionEventSource implements SourceFunction<TransactionEvent> {
        private volatile boolean running = true;
        private final Random random = new Random();
        private final String[] userIds = {"user1", "user2", "user3", "user4", "user5"};
        private int transactionIdCounter = 1;
        
        @Override
        public void run(SourceContext<TransactionEvent> ctx) throws Exception {
            while (running) {
                TransactionEvent event = new TransactionEvent(
                    userIds[random.nextInt(userIds.length)],
                    "tx_" + (transactionIdCounter++),
                    100.0 + random.nextDouble() * 900.0, // 100-1000随机金额
                    System.currentTimeMillis()
                );
                
                ctx.collect(event);
                Thread.sleep(1000 + random.nextInt(2000)); // 1-3秒随机间隔
            }
        }
        
        @Override
        public void cancel() {
            running = false;
        }
    }
    
    /**
     * 运行ValueState示例
     */
    public static void runValueStateExample(StreamExecutionEnvironment env) throws Exception {
        System.out.println("🚀 开始运行ValueState示例...");
        
        // 1. 用户登录状态跟踪示例
        DataStream<UserLoginEvent> loginEventStream = env.addSource(new MockUserLoginEventSource())
            .name("用户登录事件源");
        
        DataStream<UserStateInfo> userStateStream = loginEventStream
            .keyBy(event -> event.userId)
            .flatMap(new UserLoginStateFunction())
            .name("用户登录状态跟踪");
        
        userStateStream.print("用户状态");
        
        // 2. 交易金额累加器示例
        DataStream<TransactionEvent> transactionEventStream = env.addSource(new MockTransactionEventSource())
            .name("交易事件源");
        
        DataStream<Tuple3<String, Double, Integer>> transactionSummaryStream = transactionEventStream
            .keyBy(event -> event.userId)
            .process(new TransactionAmountAccumulator())
            .name("交易金额累加器");
        
        transactionSummaryStream.print("交易汇总");
        
        // 3. 会话超时检测示例
        DataStream<Tuple2<String, String>> sessionStatusStream = loginEventStream
            .keyBy(event -> event.userId)
            .process(new SessionTimeoutDetector())
            .name("会话超时检测");
        
        sessionStatusStream.print("会话状态");
        
        System.out.println("✅ ValueState示例配置完成");
    }
    
    /**
     * 单独运行ValueState示例的主方法
     */
    public static void main(String[] args) throws Exception {
        // 创建执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        
        // 配置状态后端
        com.flink.hbase.statemanagement.config.StateBackendConfig
            .configureDevRocksDB(env);
        
        // 运行示例
        runValueStateExample(env);
        
        // 执行作业
        env.execute("ValueState托管状态示例");
    }
} 